# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import sys
sys.path.append('../')

from PyQt4.QtGui import *
from PyQt4 import QtCore

import openbabel
import pybel

from math import atan2, pi, sqrt

from compound import Compound
from atom import Atom
from bond import Bond
from bondHelper import BondHelper

from tools import *

import common

class EditorScene(QGraphicsScene):

	def __init__(self, mol, parent=None):
		QGraphicsScene.__init__(self, parent)

		if mol:
			self.compound = Compound(mol.OBMol)
		else:
			self.compound = Compound()

		self.addItem(self.compound)
		
		self.oldSelection = []

		#self.undoStack = QUndoStack()

		self.tool = None
		self.bond = ""
		self.item = None
		self.degree = 0
		self.double = False

		self.drawingBond = False
		self.beginDrawingBondAtom = None
		self.bondEnd = None
		self.bondEndSet = False

		self.rotating = False
		self.rotationStart = None

		self.connect(self, QtCore.SIGNAL("selectionChanged()"), self.newSelection)

	def loadFile(self):
		mol = common.loadFile(self.parent())
		if mol:
			self.removeItem(self.compound)
			self.compound = Compound(mol.OBMol)
			self.addItem(self.compound)

	def saveFile(self):
		return common.saveFile(self.parent(), pybel.Molecule(self.compound.mol))

	def exportImage(self):
		self.clearSelection()
		
		image = QImage(self.sceneRect().width(), self.sceneRect().height(), QImage.Format_RGB32)
		image.fill(QColor(255,255,255).rgb())
		painter = QPainter(image)
		self.render(painter)
		return image

	def getMol(self):
		#self.compound.mol.PerceiveBondOrders()
		#self.compound.buildMol()
		return self.compound.mol

	def addMolecule(self, mol):
		self.compound.addMolecule(mol)

	def setTool(self, tool):
		self.tool = tool

	def setBond(self, bond):
		self.bond = bond

	def setRing(self, degree, double=False):
		self.item = None
		self.degree = degree
		self.double = double

	def setAtom(self, item):
		self.degree = None
		self.item = item

	def mousePressEvent(self, mouseEvent):
		if self.tool:
			self.tool(mouseEvent)
		if not mouseEvent.isAccepted():
			QGraphicsScene.mousePressEvent(self, mouseEvent)

	def mouseMoveEvent(self, mouseEvent):
		if self.drawingBond:
			self.bondEnd = self.bondHelper.changeEnd(mouseEvent.scenePos())
			self.bondEndSet = True
		elif self.rotating:
			items = self.selectedItems()
			if not items or len(items) > 1 or type(items[0]) != Compound:
				return QGraphicsScene.mouseMoveEvent(self, mouseEvent)
			rect = items[0].boundingRect()
			center = rect.center()
			v1 = self.rotationStart - center
			v2 = mouseEvent.scenePos() - center
			v1 /= sqrt(v1.x()**2+v1.y()**2)
			v2 /= sqrt(v2.x()**2+v2.y()**2)
			angle = (atan2(v1.x(), v1.y())-atan2(v2.x(), v2.y()))*180/pi
			items[0].rotate(angle)
			self.rotationStart = mouseEvent.scenePos()
		else:
			QGraphicsScene.mouseMoveEvent(self, mouseEvent)

	def mouseReleaseEvent(self, mouseEvent):
		if not self.drawingBond:
			if self.rotating:
				self.rotating = False
				self.compound.doneRotate()
				return
			QGraphicsScene.mouseReleaseEvent(self, mouseEvent)
			return
		
		self.drawingBond = False
		self.removeItem(self.bondHelper)
		endAtom = None
		
		if self.bondEndSet:
			endPos = self.bondEnd
		else:
			endPos = mouseEvent.scenePos()
		self.bondEndSet = False
		for item in self.items(endPos):
			if type(item) == Atom:
				endAtom = item
				break
		if endAtom and endAtom.atom.GetIdx() == self.beginDrawingBondAtom.atom.GetIdx():
			if type(self.item) == openbabel.OBAtom:
				endAtom.setAtomicNum(self.item.GetAtomicNum())
				#command = ChangeAtomicNum(endAtom, self.item.GetAtomicNum())
				#self.undoStack.push(command)
			return
		
		if endAtom:
			endAtom = endAtom.atom
		
		if not endAtom:
			if type(self.item) == openbabel.OBAtom:
				endAtom = self.item
			else:
				endAtom = openbabel.OBAtom()
				endAtom.SetAtomicNum(6)
			self.compound.addAtom(endAtom, self.bondEnd)
			#command = AddAtom(endAtom, self, self.bondEnd)
			#self.undoStack.push(command)
		
		endAtom = None
		for item in self.items(self.bondEnd):
			if type(item) == Atom:
				endAtom = item.atom
				break

		self.compound.addBond(self.beginDrawingBondAtom.atom, endAtom, self.bond)
		#command = AddBond(self.beginDrawingBondAtom.atom, endAtom, self.bond, self)
		#self.undoStack.push(command)
		self.beginDrawingBondAtom.updateSymbol()
		self.beginDrawingBondAtom.selected = False
		self.beginDrawingBondAtom = None

	def newSelection(self):
		for item in self.oldSelection:
			item.setSelected(False)
		self.oldSelection = self.selectedItems()
		for item in self.oldSelection:
			item.setSelected(True)

	def clear(self):
		self.oldSelection = []
		
		QGraphicsScene.clear(self)

		self.compound = Compound()

		self.addItem(self.compound)

		#self.undoStack = QUndoStack()


	def draw(self, mouseEvent):
		pos = mouseEvent.scenePos()

		items = self.items(pos.x(), pos.y(), 1, 1, QtCore.Qt.IntersectsItemBoundingRect)
		atom = None
		for item in items:
			if type(item) == Atom:
				self.beginDrawingBondAtom = item
				self.drawingBond = True
				self.bondHelper = BondHelper(item.scenePos(), self.bond)
				self.addItem(self.bondHelper)
				return
			if type(item) == Bond:
				if self.degree:
					self.compound.addRing(self.degree, bond=item.bond, pos=pos, double=self.double)
				else:
					item.setType(self.bond)
				#command = ChangeBondType(item, self.bond)
				#self.undoStack.push(command)
				return

		if self.item:
			self.compound.addAtom(self.item, pos)
			#command = AddAtom(self.item, self, pos)
			#self.undoStack.push(command)

		elif self.degree:
			self.compound.addRing(self.degree, pos=pos, double=self.double)
			#command = AddMolecule(self.item, self, pos)
			#self.undoStack.push(command)

	def move(self, mouseEvent):
		mouseEvent.ignore()

	def rotate(self, mouseEvent):
		self.rotating = True
		self.rotationStart = mouseEvent.scenePos()
		mouseEvent.accept()

	def flip(self, mouseEvent):
		items = self.items(mouseEvent.scenePos())
		if not items:
			return

		for item in items:
			if type(item) == Bond:
				#command = FlipBond(item.bond, self)
				self.compound.flipAtBond(item.bond)
				return

	def erase(self, mouseEvent):
		items = self.items(mouseEvent.scenePos())
		if not items:
			return

		toDel = None
		command = None
		for item in items:
			if type(item) == Atom:
				#toDel = item
				#command = RemoveAtom(item.atom, self)
				self.compound.removeAtom(item.atom)
				#break
				return
			elif type(item) == Bond:
				#command = RemoveBond(item.bond, self)
				#toDel = item
				self.compound.removeBond(item.bond)
				return
		
		if not toDel:
			return

		print "Found", type(toDel)

		#self.undoStack.push(command)

		mouseEvent.accept()

	def minimizeCompound(self):
		self.compound.minimize()

	def setShowHydrogens(self, shown):
		self.compound.setShowHydrogens(shown)